{-# LANGUAGE PatternGuards              #-}
{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  WumpusExpr.ConstFold
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Constant folding for K Normal Form.
-- 
-- Constant folding in action will generated redundant lets,
-- theses are removed by ElimLet.
--
--------------------------------------------------------------------------------

module WumpusExpr.ConstFold
  (
    
    constFold

  ) where


import WumpusExpr.FreshMon
import WumpusExpr.KNormal
import WumpusExpr.Operators
import WumpusExpr.Radian
import WumpusExpr.VarId
import WumpusExpr.Types

import Control.Applicative
import qualified Data.Map as M


--------------------------------------------------------------------------------
-- Constant folding

type Env = M.Map VarId Expr


findVal :: VarId -> Env -> Maybe Value
findVal vid env = M.lookup vid env >>= filt
  where
    filt (Val v) = Just v
    filt _       = Nothing

constFold :: Program -> Mon Program 
constFold (Program main_proc defs) = 
    Program <$> constFoldF main_proc <*> mapM constFoldF defs


constFoldF :: Fundef -> Mon Fundef
constFoldF (Fundef hd args body) = Fundef hd args <$> cf M.empty body

cf :: Env -> Expr -> Mon Expr
cf env (Val (Var x))  
    | Just v <- findVal x env   = return $ Val v

cf _   (Val v)                  = return $ Val v

cf env (FunCall v vs)           = FunCall v <$> mapM (cfV env) vs

cf env (Let vid e1 e2)  = do
    ea <- cf env e1
    eb <- cf (M.insert vid ea env) e2
    return $ Let vid ea eb

cf env (Ifte v e1 e2)           = go <$> cfV env v <*> cf env e1 <*> cf env e2
  where
    go (Bool True)  tc _  = tc
    go (Bool False) _  fc = fc
    go val          tc fc = Ifte val tc fc

cf env (Cmd cmd)                = Cmd <$> cfC env cmd

cf env (e1 :>> e2)              = (:>>) <$> cf env e1 <*> cf env e2
cf env (Return e)               = Return <$> cf env e
     
cfV :: Env -> Value -> Mon Value
cfV env (Var x)   | Just v <- findVal x env  = return v
cfV env (UnaryV op v1)        = go op <$> cfV env v1
  where
    go NEG      = unary NEG (negate) (negate) (negate)
    go _        = UnaryV op

cfV env (BinV op v1 v2)         = go op <$> cfV env v1 <*> cfV env v2 
  where
    go ADD      = binary ADD     (+) (+) (+)
    go SUB      = binary SUB     (-) (-) (-)
    go MUL      = binary MUL     (*) (*) (*)
    go DIV      = binary DIV     div (/) (/)
    go _        = BinV op

cfV env (RelV op v1 v2)         = go <$> cfV env v1 <*> cfV env v2 
  where
    go (Unit)       (Unit)          = Bool $ relOrd op () ()
    go (Bool s)     (Bool t)        = Bool $ relOrd op s t
    go (Int s)      (Int t)         = Bool $ relOrd op s t
    go (Float s)    (Float t)       = Bool $ relOrd op s t
    go (Angle s)    (Angle t)       = Bool $ relOrd op s t
    go (String s)   (String t)      = Bool $ relOrd op s t
    go (PsSymbol s) (PsSymbol t)    = Bool $ relOrd op s t
    go va           vb              = RelV op va vb

cfV _   v                       = return v


cfC :: Env -> Command -> Mon Command
cfC _   (NewPath)                   = pure NewPath
cfC _   (ClosePath)                 = pure ClosePath
cfC _   (Fill)                      = pure Fill
cfC _   (Stroke)                    = pure Stroke
cfC env (MoveTo x y)                = MoveTo <$> cfV env x <*> cfV env y
cfC env (LineTo x y)                = LineTo <$> cfV env x <*> cfV env y

cfC env (CurveTo x1 y1 x2 y2 x3 y3) = 
    CurveTo <$> cfV env x1 <*> cfV env y1 
            <*> cfV env x2 <*> cfV env y2
            <*> cfV env x3 <*> cfV env y3

cfC _   (GSave)                     = pure GSave
cfC _   (GRestore)                  = pure GRestore
cfC env (Scale x y)                 = Scale <$> cfV env x <*> cfV env y
cfC env (Rotate ang)                = Rotate <$> cfV env ang

cfC env (Translate x y)             = Translate <$> cfV env x <*> cfV env y
cfC env (Transform a b c d e f)     =    
    Transform <$> cfV env a <*> cfV env b <*> cfV env c
              <*> cfV env d <*> cfV env e <*> cfV env f

cfC env (ShowText v)                = ShowText <$> cfV env v


--------------------------------------------------------------------------------
-- Helpers

unary :: UnaryOp -> (Int -> Int) -> (Decimal -> Decimal) -> (Radian -> Radian)
      -> Value -> Value 
unary _  opi _   _   (Int i)    = Int   $ opi i
unary _  _   opf _   (Float i)  = Float $ opf i
unary _  _   _   opr (Angle i)  = Angle $ opr i
unary op _   _   _   v1         = UnaryV op v1


binary :: BinOp 
       -> (Int -> Int -> Int) 
       -> (Decimal -> Decimal -> Decimal)
       -> (Radian -> Radian -> Radian)
       -> Value -> Value -> Value
binary _  opi _   _   (Int i)   (Int j)    = Int   $ i `opi` j
binary _  _   opf _   (Float i) (Float j)  = Float $ i `opf` j
binary _  _   _   opr (Angle i) (Angle j)  = Angle $ i `opr` j
binary op _   _   _   v1        v2         = BinV op v1 v2



relOrd :: Ord a => RelOp -> a -> a -> Bool
relOrd EQU     a b = a == b
relOrd NOT_EQU a b = a /= b
relOrd LTHAN   a b = b <  a
relOrd GTHAN   a b = a >  b
relOrd LT_EQU  a b = a <= b
relOrd GT_EQU  a b = a >= b

